"""
软件名：PlotDub
功能：本文件用于启动剧情配音软件 PlotDub WebUI
作者：hxp@海小皮
版本： V1.1
功能：
* 加载模型环境参数
* 启动 WebUIUI 界面
软件代码及相关文件（夹）：
* roles.json 游戏角色模型权重、参考音频路径配置文件
* inference_param.json 模型推理参数配置文件
* PlotDub_weights/ 游戏角色模型权重、参考音频文件夹
* TEMP/ 软件运行时的临时数据缓存
* PlotDub/ 软件模块代码
"""


"""******************************* PlotDub 软件初始化 ************************************"""
import os, sys, shutil
# 检查/创建 TEMP 文件夹
now_dir = os.getcwd()
sys.path.insert(0, now_dir)
tmp = os.path.join(now_dir, "TEMP")
os.makedirs(tmp, exist_ok=True)
os.environ["TEMP"] = tmp
if os.path.exists(tmp):
    for name in os.listdir(tmp):
        if name == "jieba.cache": continue
        path = "%s/%s"%(tmp, name)
        delete = os.remove if os.path.isfile(path) else shutil.rmtree
        try:
            delete(path)
        except Exception as e:
            print(str(e))
            pass

from PlotDub.DataStruct import GameType, ScreenType, EmotionType, Emotion, VersionType
# 软件版本参数
plot_dub_version = "V1.1"
# 角色信息 json 文件名
roles_json_file_name = "roles.json"
# 推理参数 json 文件名
inference_json_file_name = "inference_param.json"
# 权重文件存放文件夹名
weights_folder_name = "PlotDub_weights"
# WebUI 顶部显示软件名
plot_dub_name = f"原神&星铁剧情配音软件 {plot_dub_version}"
# WebUI 顶部软件简要说明
plot_dub_illustration = (
    f"软件作者：hxp@海小皮<br>"
    f"软件版本：{plot_dub_version}<br>"
    "本软件基于 GPT-SoVITS 开发. <br>"
    "本软件继承 GPT-SoVITS 开源协议."
)
# WebUI 文件长传格式提示字段
file_update_note = (
    "GPT 模型权重文件后缀 .ckpt<br>"
    "SoVITS 模型权重文件后缀 .pth<br>"
    "参考音频上传注意:<br>"
    "---1. 格式为 情感-参考文本.wav<br>"
    f"---2. 情感有效字段:{EmotionType.CALM.value} {EmotionType.HAPPY.value} {EmotionType.ANGRY.value} {EmotionType.QUESTIONING.value}<br>"
    "---3. 长度为 3-10 秒"
)  # TODO: EmotionType 修改后，这里也需要修改
# 软件支持的游戏列表
games_list = [GameType.GENSHIN_IMPACT.value,
              GameType.STAR_RAIL.value]

# 使用的 GPT-SoVITS 模型版本
model_version = VersionType.V2

# 设置推理语言
# GPT-SoVITS v1 可选 中文 英文 日文 中英混合 日英混合 多语种混合
# GPT-SoVITS v2 可选 中文 英文 日文 粤语 韩文 中英混合 日英混合 粤英混合 韩英混合 多语种混合 多语种混合(粤语)
# 目前只有中文能用，其他貌似需要下载模型
language = "中文"
os.environ["language"] = language

import json
# 检查/创建 {weights_folder_name}/gameName 文件夹
for folder_name in games_list:
    os.makedirs(f"{weights_folder_name}/{folder_name}",exist_ok=True)

# 检查/创建默认 {roles_json_file_name}
if os.path.exists(f"./{roles_json_file_name}"):
    pass
else:
    with open(f"./{roles_json_file_name}", 'w', encoding="utf-8") as file:
        json.dump({'原神': {}, '星铁': {}}, file)

# 检查/创建默认 {inference_json_file_name} 文件
if os.path.exists(f"./{inference_json_file_name}"):
    pass
else:
    with open(f"./{inference_json_file_name}", 'w', encoding="utf-8") as file:
        json.dump({'if_freeze': False,
                   'speed': 1,
                   'top_k': 15,
                   'top_p': 1,
                   'temperature': 1}, file)

# 加载 roles_json_file_name 文件，初始化 game_roles_dict game_roles_list_dict role_emotions_list_dict
from PlotDub.hxpTTS import Role
games_roles_dict = {}  # 用于初始化 ITS 对象 key = game + role_name value = Role 对象
game_roles_list_dict = {}  # 用于 WebUI 角色列表的显示查表 key = GameType.value value = [角色名字列表]
role_emotions_list_dict = {}  # 用于 WebUI 情感列表的显示查表 key = 角色名 value = [角色情感列表 EmotionType.value]

with open(f"./roles.json", 'r', encoding="utf-8") as file:
    roles_data = json.load(file)
    # 解析 json 文件
    for game, roles in roles_data.items():
        roles_list = []
        # 解析角色信息，构建 Role 对象
        for role_name, inform in roles.items():
            emotion_dict = {}
            emotions_list = []
            gpt_path = inform.get("GPT", "")
            sovits_path = inform.get("SoVITS", "")

            # 解析情感参考音频和参考文本，构建 emotion_dict
            emotions = inform.get("Emotion", {})
            for emo, ref_wav_path in emotions.items():
                emotion_dict[EmotionType(emo)] = Emotion(ref_wav_path)
                emotions_list.append(emo)

            # 填充 game_roles_dict
            games_roles_dict[game + role_name] = Role(GameType(game), role_name, gpt_path, sovits_path, emotion_dict)

            # 填充 role_emotions_list_dict
            role_emotions_list_dict[role_name] = emotions_list

            # 填充 roles_list
            roles_list.append(role_name)

        # 填充 game_roles_list_dict
        game_roles_list_dict[game] = roles_list

# 加载 {inference_json_file_name} 文件，初始化 inference_param
from PlotDub.hxpTTS import InferenceParam
inference_param = InferenceParam(language, model_version)
inference_param.LoadFromJson(f"./{inference_json_file_name}")

# 创建后端对象
from PlotDub.hxpITS import hxpITS
its = hxpITS(games_roles_dict, inference_param)

"""**************************************************************************************"""

"""************************************ WebUI 代码 ***************************************"""
# WebUI
try:
    import gradio.analytics as analytics

    analytics.version_check = lambda: None
except:
    ...
import gradio as gr
from PlotDub.hxpWebUI import html_left_title, html_all_center_title, html_left_text


infer_ttswebui = 9872
infer_ttswebui = int(infer_ttswebui)
is_share = "False"
is_share = eval(is_share)


def game_try_dropdown_change_callback(game):
    """
    游戏名下拉菜单修改回调函数
    更新 roles_list
    """
    new_roles_list = game_roles_list_dict.get(game)
    if new_roles_list:
        return {"__type__": "update", "choices": new_roles_list, "value": new_roles_list[0]}
    else:
        return {"__type__": "update", "choices": [], "value": "None"}


def role_dropdown_change_callback(role_name):
    """
    角色名下拉菜单修改回调函数
    更新 emotions_list
    """
    new_emotions_list = role_emotions_list_dict.get(role_name)
    if new_emotions_list:
        return {"__type__": "update", "choices": new_emotions_list,"value": new_emotions_list[0]}
    else:
        return {"__type__": "update", "choices": [], "value": "None"}


def try_listen_button_callback(game, role_name, eno):
    """
    试听按钮回调函数
    """
    print(f"************************【WebUI】***************************")
    print(f"正在试听: 【{game}】【{role_name}】情绪为【{eno}】的参考音频")
    print(f"***********************************************************")
    its.TryListen(GameType(game), role_name, EmotionType(eno))


def merge_button_callback(game, role_name, emo, text):
    """
    合成按钮回调函数
    """
    print(f"************************【WebUI】***************************")
    print(f"正在生成: 【{game}】【{role_name}】情绪为【{emo}】内容为【{text}】音频")
    print(f"***********************************************************")
    audio = its.TryMerge(GameType(game), role_name, EmotionType(emo), text)
    return audio


def add_role_button_callback(game, role_name, gpt_weight, sovits_weight, ref_wavs):
    """
    添加角色按钮回调函数
    """
    # 校验文件后缀
    if not gpt_weight.endswith('.ckpt'):
        print(f"************************【WebUI】***************************")
        print(f"错误: GPT 模型权重文件 {gpt_weight} 的后缀必须是 .ckpt")
        print(f"【{game}】: 角色【{role_name}】添加失败")
        print(f"***********************************************************")
        return
    if not sovits_weight.endswith('.pth'):
        print(f"************************【WebUI】***************************")
        print(f"错误: SoVITS 模型权重文件 {sovits_weight} 的后缀必须是 .pth")
        print(f"【{game}】: 角色【{role_name}】添加失败")
        print(f"***********************************************************")
        return
    for ref_wav in ref_wavs:
        if not ref_wav.endswith('.wav'):
            print(f"************************【WebUI】***************************")
            print(f"错误: 参考音频文件 {ref_wav} 的后缀必须是 .wav")
            print(f"【{game}】: 角色【{role_name}】添加失败")
            print(f"***********************************************************")
            return

    print(f"************************【WebUI】***************************")
    print(f"正在添加【{game}】: 角色【{role_name}】的模型权重和参考音频 ")
    print(f"GPT 模型权重文件: {gpt_weight}")
    print(f"SoVITS 模型权重文件: {sovits_weight}")
    print(f"参考音频文件 : {ref_wavs}")

    role_path = f"{weights_folder_name}/{game}/{role_name}"

    # 加载现有 JSON 数据
    if os.path.exists(roles_json_file_name):
        with open(roles_json_file_name, 'r', encoding='utf-8') as f:
            roles_data = json.load(f)
    else:
        roles_data = {}

    # 检查角色是否已经存在
    if os.path.exists(role_path):
        # 存在则更新文件并更新 json 文件
        print(f"{role_path} 已存在，正在更新文件...")
        # 清空角色文件夹
        shutil.rmtree(role_path)
        os.makedirs(role_path)
    else:
        # 不存在则创建文件夹并拷贝文件 同时在 json 中添加新的条目
        print(f"{role_path} 不存在，正在创建文件夹...")
        # 创建文件见
        os.makedirs(role_path)

    # 拷贝文件到角色目录
    shutil.copy(gpt_weight, os.path.join(role_path, os.path.basename(gpt_weight)))
    shutil.copy(sovits_weight, os.path.join(role_path, os.path.basename(sovits_weight)))
    for ref_wav in ref_wavs:
        shutil.copy(ref_wav, os.path.join(role_path, os.path.basename(ref_wav)))

    print(f"文件保存成功.")

    # 更新 JSON 数据
    emotions_dict = {
        Emotion(ref_wav).type.value: os.path.join(role_path, os.path.basename(ref_wav))
        for ref_wav in ref_wavs
    }
    roles_data.setdefault(game, {})[role_name] = {
        "GPT": os.path.join(role_path, os.path.basename(gpt_weight)),
        "SoVITS": os.path.join(role_path, os.path.basename(sovits_weight)),
        "Emotion": emotions_dict
    }

    # 保存 JSON 文件
    with open(roles_json_file_name, 'w', encoding='utf-8') as f:
        json.dump(roles_data, f, ensure_ascii=False, indent=4)

    print(f"{roles_json_file_name} 文件更新完成.")

    # 更新 WebUI 下拉菜单字典
    if role_name not in game_roles_list_dict[game]:
        game_roles_list_dict[game].append(role_name)  # 更新游戏角色字典
    role_emotions_list_dict[role_name] = [e for e in emotions_dict]  # 更新角色情感字典

    print(f"网页下拉列表已更新.")

    print(f"***********************************************************")

    # 更新 ITS
    new_gpt_path = roles_data[game][role_name]["GPT"]  # 新的 gpt_weight 文件路径
    new_sovits_path = roles_data[game][role_name]["SoVITS"]  # 新的 sovits_weight 文件路径
    new_emotions_dict = {
        EmotionType(emo): Emotion(ref_wav)
        for emo, ref_wav in emotions_dict.items()
    }
    new_role = Role(GameType(game), role_name, new_gpt_path, new_sovits_path, new_emotions_dict)
    its.AddGameRole(new_role)


def screen_resolution_save_button_callback(screen_resolution):
    """
    分辨率保存按钮回调函数
    """
    print(f"************************【WebUI】***************************")
    print(f"屏幕分辨率参数【{screen_resolution}】已保存。")
    print(f"***********************************************************")
    its.SaveScreenResolution(ScreenType(screen_resolution))


# 重置按钮的回调函数
def recover_if_freeze():
    return False


def recover_speed():
    return 1


def recover_top_k():
    return 15


def recover_top_p():
    return 1


def recover_temperature():
    return 1


def inference_param_save_button_callback(if_freeze, speed, top_k, top_p, temperature):
    """
    推理参数保存按钮回调函数
    """
    print(f"************************【WebUI】***************************")
    print(f"GPT-SoVITS 模型推理参数参数已保存: ")
    print(f"if_freeze: {if_freeze}")
    print(f"speed: {speed}")
    print(f"top_k: {top_k}")
    print(f"top_p: {top_p}")
    print(f"temperature: {temperature}")

    # 加载现有 JSON 数据
    if os.path.exists(inference_json_file_name):
        with open(inference_json_file_name, 'r', encoding='utf-8') as f:
            inference_param_data = json.load(f)
    else:
        inference_param_data = {}

    # 保存新的参数
    inference_param_data['if_freeze'] = if_freeze
    inference_param_data['speed'] = speed
    inference_param_data['top_k'] = top_k
    inference_param_data['top_p'] = top_p
    inference_param_data['temperature'] = temperature

    # 保存 JSON 文件
    with open(inference_json_file_name, 'w', encoding='utf-8') as f:
        json.dump(inference_param_data, f, ensure_ascii=False, indent=4)

    print(f"{inference_json_file_name} 文件更新完成.")

    print(f"***********************************************************")

    its.SaveInferenceParam(InferenceParam(if_freeze=if_freeze, speed=speed, top_k=top_k, top_p=top_p, temperature=temperature))


def run_button_callback(game):
    """
    运行按钮回调函数
    """
    print(f"************************【WebUI】***************************")
    print(f"【WebUI】【{game}】剧情配音开始")
    print(f"***********************************************************")
    its.Run(GameType(game))  # 确保传入的是 GameType 枚举类对象


def stop_button_callback():
    """
    停止按钮回调函数
    """
    print(f"************************【WebUI】***************************")
    print(f"【WebUI】剧情配音停止")
    print(f"***********************************************************")
    its.Stop()


def WebUI():
    # 初始化下拉菜单
    # 根据 GameType 初始化 games_list
    games_list = [GameType.GENSHIN_IMPACT.value,
                  GameType.STAR_RAIL.value]
    # 根据 game_roles_list_dict 初始化 roles_list
    roles_list = game_roles_list_dict[GameType.GENSHIN_IMPACT.value]
    # 根据 role_emotions_list_dict 初始化 emotions_list
    emotions_list = role_emotions_list_dict[roles_list[0]]
    # 支持的分辨率列表
    resolutions_list = [ScreenType.S_1080P.value,
                        ScreenType.S_2K.value,
                        ScreenType.S_4K.value]

    with gr.Blocks(title="PlotDub WebUI") as app:
        gr.Markdown(html_all_center_title(plot_dub_name, 'h1'))
        # 软件说明
        gr.Markdown(value=plot_dub_illustration)
        # 操作界面
        with gr.Group():
            # 当前角色列表试听
            gr.Markdown(html_all_center_title("当前角色列表试听/合成", 'h2'))
            with gr.Row():
                # 该部分竖向排布
                with gr.Column():
                    # 第一行选择游戏、角色、情感+试听按钮
                    with gr.Row():
                        # 游戏选择框
                        Game_try_dropdown = gr.Dropdown(
                            label="游戏选择",
                            choices=games_list,
                            value=games_list[0],
                            interactive=True,
                            scale=30
                        )
                        # 角色选择框
                        Role_dropdown = gr.Dropdown(
                            label="角色选择",
                            choices=roles_list,  # 按拼音排序
                            value=roles_list[0],
                            interactive=True,
                            scale=30
                        )
                        # 情感选择框
                        Emo_dropdown = gr.Dropdown(
                            label="情感选择",
                            choices=emotions_list,
                            value=emotions_list[0],
                            interactive=True,
                            scale=30
                        )
                        Game_try_dropdown.change(
                            game_try_dropdown_change_callback,
                            [Game_try_dropdown],
                            [Role_dropdown]
                        )
                        Role_dropdown.change(
                            role_dropdown_change_callback,
                            [Role_dropdown],
                            [Emo_dropdown]
                        )
                        try_listen_button = gr.Button("试听", variant="primary", scale=10)
                        try_listen_button.click(
                            try_listen_button_callback,
                            [Game_try_dropdown, Role_dropdown, Emo_dropdown],
                            []
                        )
                    # 第二行，输入文本框和生成按钮
                    with gr.Row():
                        target_text = gr.Textbox(label="目标文本", value="", lines=1, max_lines=1, scale=90)
                        merge_button = gr.Button("生成", variant="primary", scale=10)
                    # 第三行，音频播放器
                    merge_result = gr.Audio(label="生成语音", interactive=False)
                    merge_button.click(merge_button_callback, [Game_try_dropdown, Role_dropdown, Emo_dropdown, target_text], [merge_result])
            # 添加角色
            gr.Markdown(html_all_center_title("添加角色", 'h2'))
            with gr.Row():
                with gr.Column(scale=30):
                    Game_add_dropdown = gr.Dropdown(
                        label="游戏",
                        choices=games_list,
                        value=games_list[0],
                        interactive=True,
                        scale=30
                    )
                    role_name = gr.Textbox(label="角色名字", value="", lines=1, max_lines=1, scale=30)
                    gr.Markdown(html_all_center_title("注意事项", label='h3'))
                    gr.Markdown(html_left_text(file_update_note, size=15))
                gpt_weight_file = gr.File(label="上传 GPT .ckpt 权重文件", scale=20)
                sovits_weight_file = gr.File(label="上传 SoVITS .pth 权重文件", scale=20)
                ref_wavs = gr.Files(label="上传 角色参考音频 文件", scale=20)
                add_role_button = gr.Button("添加", variant="primary", scale=10)
                add_role_button.click(
                    add_role_button_callback,
                    [Game_add_dropdown, role_name, gpt_weight_file, sovits_weight_file, ref_wavs]
                )
            # 屏幕分辨率设置
            gr.Markdown(html_all_center_title("屏幕分辨率设置", 'h2'))
            with gr.Row():
                Resolution_dropdown = gr.Dropdown(
                    label="屏幕分辨率选择",
                    choices=resolutions_list,
                    value=resolutions_list[0],
                    interactive=True,
                    scale=90
                )
                screen_resolution_save_button = gr.Button("保存", variant="primary", scale=10)
                screen_resolution_save_button.click(
                    screen_resolution_save_button_callback,
                    [Resolution_dropdown]
                )
            # 推理参数设置
            gr.Markdown(html_all_center_title("GPT-SoVITS 推理参数设置", 'h2'))
            with gr.Row():
                with gr.Column(scale=90):
                    with gr.Row():
                        if_freeze = gr.Checkbox(label="是否直接对上次合成结果调整语速和音色。防止随机性。", value=inference_param.if_freeze,
                                                interactive=True, show_label=True, scale=9)
                        if_freeze_recover_button = gr.Button("恢复默认值", scale=1)
                        if_freeze_recover_button.click(recover_if_freeze, [], [if_freeze])
                    with gr.Row():
                        speed = gr.Slider(minimum=0.6, maximum=1.65, step=0.05, label="语速", value=inference_param.speed, interactive=True,
                                          scale=9)
                        speed_recover_button = gr.Button("恢复默认值", scale=1)
                        speed_recover_button.click(recover_speed, [], [speed])
                    gr.Markdown(html_all_center_title("GPT采样参数(无参考文本时不要太低。不懂就用默认)："))
                    with gr.Row():
                        top_k = gr.Slider(minimum=1, maximum=100, step=1, label="top_k", value=inference_param.top_k, interactive=True,
                                          scale=9)
                        top_k_recover_button = gr.Button("恢复默认值", scale=1)
                        top_k_recover_button.click(recover_top_k, [], [top_k])
                    with gr.Row():
                        top_p = gr.Slider(minimum=0, maximum=1, step=0.05, label="top_p", value=inference_param.top_p, interactive=True,
                                          scale=9)
                        top_p_recover_button = gr.Button("恢复默认值", scale=1)
                        top_p_recover_button.click(recover_top_p, [], [top_p])
                    with gr.Row():
                        temperature = gr.Slider(minimum=0, maximum=1, step=0.05, label="temperature", value=inference_param.temperature,
                                                interactive=True, scale=9)
                        temperature_recover_button = gr.Button("恢复默认值", scale=1)
                        temperature_recover_button.click(recover_temperature, [], [temperature])
                inference_param_save_button = gr.Button("保存", variant="primary", scale=10)
                inference_param_save_button.click(
                    inference_param_save_button_callback,
                    [if_freeze, speed, top_k, top_p, temperature]
                )
            # 启停按钮
            gr.Markdown(html_all_center_title("剧情配音", 'h2'))
            with gr.Row():
                Game_run_dropdown = gr.Dropdown(
                    label="游戏",
                    choices=games_list,
                    value=games_list[0],
                    interactive=True,
                    scale=15
                )
                run_button = gr.Button("启动", variant="primary", scale=50)
                run_button.click(run_button_callback, [Game_run_dropdown])
                stop_button = gr.Button("停止", variant="stop", scale=50)
                stop_button.click(stop_button_callback)

    app.queue().launch(  # concurrency_count=511, max_size=1022
        server_name="0.0.0.0",
        inbrowser=True,
        share=is_share,
        server_port=infer_ttswebui,
        quiet=True,
    )


"""**************************************************************************************"""


if __name__ == '__main__':
    WebUI()
